import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import json
import sys
import os
import warnings
import math
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def replace_invalid_floats(data):
    """Replace NaN, inf, and -inf with None."""
    if isinstance(data, dict):
        return {key: replace_invalid_floats(value) for key, value in data.items()}
    elif isinstance(data, list):
        return [replace_invalid_floats(item) for item in data]
    elif isinstance(data, float):
        if math.isnan(data) or math.isinf(data):
            return None
    return data

def main():
    try:
        if len(sys.argv) != 3:
            result = {
                'status': 'error',
                'message': '参数错误，需要提供CSV文件路径和选择的特征JSON文件路径'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        csv_file_path = sys.argv[1]
        features_json_file = sys.argv[2]

        # 验证输入文件是否存在
        if not os.path.exists(csv_file_path):
            result = {
                'status': 'error',
                'message': f'CSV文件不存在: {csv_file_path}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        if not os.path.exists(features_json_file):
            result = {
                'status': 'error',
                'message': f'特征JSON文件不存在: {features_json_file}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 读取特征列表
        try:
            with open(features_json_file, 'r', encoding='utf-8') as f:
                selected_features = json.load(f)

            if not isinstance(selected_features, list):
                result = {
                    'status': 'error',
                    'message': f'特征JSON文件格式错误，应为列表格式'
                }
                print(json.dumps(result, ensure_ascii=False))
                return

        except json.JSONDecodeError as e:
            result = {
                'status': 'error',
                'message': f'特征JSON文件解析失败: {str(e)}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 读取数据
        df = pd.read_csv(csv_file_path)

        # 确保所有特征字段都存在
        missing_features = [f for f in selected_features if f not in df.columns]
        if missing_features:
            result = {
                'status': 'error',
                'message': f'数据中缺少字段: {missing_features}'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 选择20个字段的数据（包含qyb）
        selected_df = df[selected_features].copy()

        # 检查数据有效性
        if selected_df.dropna().empty:
            result = {
                'status': 'error',
                'message': '选定的特征数据全部为缺失值'
            }
            print(json.dumps(result, ensure_ascii=False))
            return

        # 计算相关性矩阵
        correlation_matrix = selected_df.corr()

        # 处理无效值
        corr_for_json = correlation_matrix.where(pd.notnull(correlation_matrix), None)
        corr_for_json = corr_for_json.applymap(
            lambda x: None if (isinstance(x, float) and (math.isnan(x) or math.isinf(x))) else x
        )

        # 转换为字典
        corr_dict = corr_for_json.to_dict()

        # 创建热力图
        heatmap_path = None
        try:
            plt.figure(figsize=(16, 14))

            # 生成热力图
            sns.heatmap(correlation_matrix,
                        annot=True,
                        cmap='RdYlBu_r',
                        center=0,
                        square=True,
                        fmt='.3f',
                        cbar_kws={"shrink": .8},
                        linewidths=0.5)

            plt.title(f'{len(selected_features)}个重要特征相关性热力图\n(包含QYB目标变量)',
                      fontsize=16, fontweight='bold', pad=20)
            plt.xticks(rotation=45, ha='right')
            plt.yticks(rotation=0)
            plt.tight_layout()

            # 确保输出目录存在
            os.makedirs('temp/output', exist_ok=True)
            heatmap_path = 'temp/output/feature_correlation_heatmap.png'
            plt.savefig(heatmap_path, dpi=300, bbox_inches='tight')
            plt.close()

            print(f"热力图已保存: {heatmap_path}")

        except Exception as e:
            print(f"热力图生成失败: {str(e)}", file=sys.stderr)

        # 分析强相关性特征对
        strong_correlations = []
        for i in range(len(correlation_matrix.columns)):
            for j in range(i+1, len(correlation_matrix.columns)):
                corr_value = correlation_matrix.iloc[i, j]
                if not (math.isnan(corr_value) or math.isinf(corr_value)) and abs(corr_value) > 0.7:
                    strong_correlations.append({
                        'feature1': correlation_matrix.columns[i],
                        'feature2': correlation_matrix.columns[j],
                        'correlation': float(corr_value)
                    })

        # 获取与qyb的相关性（如果qyb在特征中）
        qyb_correlations = {}
        if 'qyb' in corr_for_json.columns:
            for col in corr_for_json.columns:
                if col != 'qyb':
                    qyb_correlations[col] = corr_for_json.loc['qyb', col]

        # 构建结果
        result = {
            'status': 'success',
            'message': '第二次相关性分析完成 - 20个特征热力图生成',
            'analysis_type': 'feature_correlation_heatmap',
            'correlation_matrix': corr_dict,
            'features_analyzed': selected_features,
            'matrix_shape': list(correlation_matrix.shape),
            'strong_correlations': strong_correlations,
            'strong_correlation_count': len(strong_correlations),
            'qyb_correlations': qyb_correlations,
            'data_info': {
                'total_rows': len(df),
                'valid_rows': len(selected_df.dropna()),
                'features_count': len(selected_features)
            }
        }

        if heatmap_path:
            result['heatmap_path'] = heatmap_path

        # 过滤无效值并输出
        result = replace_invalid_floats(result)
        print(json.dumps(result, ensure_ascii=False))

    except Exception as e:
        error_result = {
            'status': 'error',
            'message': f'第二次相关性分析失败: {str(e)}'
        }
        print(json.dumps(error_result, ensure_ascii=False))

if __name__ == "__main__":
    main()
